import sys
sys.path.append('..')

from dataflow.DataFlow import NodeBase
from dataflow.DataFlow import DataFlowExecutor
import cv2
import matplotlib.pyplot as plt
import time
import os.path
import matlab
import matlab.engine

import dip.basic
import dip.enhance
import dip.transform
import dip.merge_seg
import dip.restoration

def get_file_save_path():
    filename = str(int(time.time()*1e6)) + '.png'
    filepath = os.path.join(os.path.dirname(__file__), "cache", filename)
    return filepath

''' 
Basic 
'''
class ImageImport(NodeBase):
    '''导入图片'''
    def __init__(self):
        inputs = [None]     #[image]
        outputs = [None]    #[image]

        super().__init__(inputs, outputs)

    def func(self):
        path = self.options["path"]
        print("@@@import image:", path)

        self.outputs[0] = cv2.imread(path, cv2.IMREAD_UNCHANGED)


class ImageDisplay(NodeBase):
    '''显示图片'''
    def __init__(self):
        inputs = [None] # [image]
        outputs = []
        super().__init__(inputs, outputs)

    def func(self):
        # 文件名及其路径
        filepath = get_file_save_path()

        # 保存图片
        cv2.imwrite(filepath, self.inputs[0])

        #添加结果
        imageFlowExecutor.add_result({
            "type":"imageShow",
            "data":{
                "path":filepath
            }
        })

        print("@@@display image:", filepath)


    def update_input(self, k, input):
        if(k == 0):
            self.inputs[k] = input.copy()   #接收到图片后先复制一份


'''
Transform
'''
class ImageGrayScale(NodeBase):
    '''灰度化'''
    def __init__(self):
        inputs = [None]
        outputs = [None]
        super().__init__(inputs, outputs)
    
    def func(self):
        self.outputs[0] = dip.transform.grayscale(self.inputs[0])
        

class ImageFFT(NodeBase):
    '''FFT'''
    def __init__(self):
        inputs = [None]
        outputs = [None]
        super().__init__(inputs, outputs)
    
    def func(self):
        spectrum = dip.transform.FFT(self.inputs[0])
        self.outputs[0] = spectrum

        filepath = get_file_save_path()
        plt.figure()
        plt.title('FFT')
        plt.imshow(spectrum, 'gray')
        plt.savefig(filepath)
        imageFlowExecutor.add_result({
            "type":"imageShow",
            "data":{
                "path":filepath
            }
        })

        

class ImageDCT(NodeBase):
    '''DCT'''
    def __init__(self):
        inputs = [None]
        outputs = [None]
        super().__init__(inputs, outputs)

    def func(self):
        spectrum = dip.transform.DCT(self.inputs[0])
        self.outputs[0] = spectrum

        filepath = get_file_save_path()
        plt.figure()
        plt.title('DCT')
        plt.imshow(spectrum, 'gray')
        plt.savefig(filepath)
        imageFlowExecutor.add_result({
            "type":"imageShow",
            "data":{
                "path":filepath
            }
        })



'''
Enhance
'''
class ImageSharpen(NodeBase):
    def __init__(self):
        inputs = [None]
        outputs = [None]
        super().__init__(inputs, outputs)

    def func(self):
        self.outputs[0] = dip.enhance.Sharpen.laplace_mask(
            self.inputs[0],
            self.options['field'],
            self.options["alpha"]
        )

class ImageLinearEnhancement(NodeBase):
    '''线性拉伸'''
    def __init__(self):
        inputs = [None]
        outputs = [None]
        
        self.operationMap = {
            "up_down": dip.enhance.LinearEnhancement.moveUpDown,
            "reverse": dip.enhance.LinearEnhancement.reverse,
            "adjustContrast" : dip.enhance.LinearEnhancement.adjustGrayContrast
        }
        super().__init__(inputs, outputs)

    def func(self):
        operation = self.operationMap[self.options['operation']]

        self.outputs[0] = operation(
            self.inputs[0],
            self.options['value']
        )

class ImageNonLinearEnhancement(NodeBase):
    def __init__(self):
        inputs = [None]
        outputs = [None]
        super().__init__(inputs, outputs)

    def func(self):
        self.outputs[0] = dip.enhance.NonLinearEnhancement.logarithm(
            self.inputs[0],
            self.options["value"]
        )

        if(self.options["showCurve"]):
            filepath = dip.enhance.NonLinearEnhancement.get_logarithm_curve(
                self.options["value"]
            )

            filepath = os.path.join(os.path.dirname(__file__), filepath)
            print("@@curve",filepath)
            imageFlowExecutor.add_result({
                "type":"imageShow",
                "data":{
                    "path":filepath
                }
            })

class ImageHistgram(NodeBase):
    def __init__(self):
        inputs = [None]
        outputs = [None]
        super().__init__(inputs, outputs)
    
    def func(self):
        plt.figure()
        img = self.inputs[0]

        plt.hist(img.ravel(), 256)
        path = get_file_save_path()
        plt.savefig(path)

        imageFlowExecutor.add_result({
            "type":"imageShow",
            "data":{
                "path": path
            }
        })

       

class ImageEqualizeHist(NodeBase):
    def __init__(self):
        inputs = [None]
        outputs = [None]
        super().__init__(inputs, outputs)
    
    def func(self):
        self.outputs[0] = dip.enhance.equalizeHist(self.inputs[0])


''' 
Merge and segment 
'''
class ImageHumanSeg(NodeBase):
    '''提取人像'''
    def __init__(self):
        inputs=[None]   #[image]
        outputs=[None]  #[image]
        super().__init__(inputs, outputs)
    
    def func(self):
        img = self.inputs[0]

        result = dip.merge_seg.hunman_seg(img)
        print("@@@image hunmanseg:", result)

        self.outputs[0] = cv2.imread(result, cv2.IMREAD_UNCHANGED)

    def update_input(self, k, input):
        if(k == 0):
            self.inputs[k] = input.copy()   #接收到图片后先复制一份




''''
图像恢复
'''
class ImageProjection(NodeBase):
    def __init__(self):
        inputs =  [None]
        outputs = [None] 
        super().__init__(inputs, outputs)

    def func(self):
        self.outputs[0] = dip.restoration.projection(self.inputs[0])


class ImageDeBlur(NodeBase):
    def __init__(self):
        inputs = [None]
        outputs = [None]
        super().__init__(inputs, outputs)

    def func(self):
        img = self.inputs[0]
        path = get_file_save_path()
        cv2.imwrite(path, img)
        
        len = self.options["len"]
        theta = self.options["theta"]
        k = self.options["k"]
        
        eng = matlab.engine.start_matlab()

        path = eng.deblur(path, float(len), float(theta), float(k))

        img_out = cv2.imread(path, cv2.IMREAD_UNCHANGED)

        self.outputs[0] = img_out

        eng.exit()


class ImageFlowExecutor(DataFlowExecutor):
    '''节点处理器'''
    def __init__(self):
        super().__init__()
        
        self._results = []

        #在这里注册节点
        self.node_type = {
            "imageImport": ImageImport,
            "imageDisplay": ImageDisplay,

            "imageHistgram" : ImageHistgram,
            "imageGrayScale": ImageGrayScale,
            "imageFFT" : ImageFFT,
            "imageDCT" : ImageDCT,

            "imageSharpen": ImageSharpen,
            "imageLinearEnhancement" : ImageLinearEnhancement,
            "imageNonLinearEnhancement" : ImageNonLinearEnhancement,
            "imageEqualizeHist" : ImageEqualizeHist,

            
            "imageHumanSeg": ImageHumanSeg,

            "imageProjection": ImageProjection,
            "imageDeBlur": ImageDeBlur
        }

    def add_result(self, result):
        self._results.append(result)

    def get_result(self):
        result = self._results.copy()
        self._results.clear()

        return result

    def load_from_json(self, data):
        nodes: list[NodeBase] = []
        for node_data in data:
            # 添加node实例
            instance = self.node_type[node_data["type"]]
            nd = instance()
            nd.id = node_data['id']
            nodes.append(nd)
        
        for n, node_data in enumerate(data):
            # front
            for j in node_data["front"]:
                nodes[n].add_front(nodes[j])
            # next
            for j in node_data["next"]:
                nodes[n].add_next(nodes[j])
            
            # out_flow
            if(node_data["out_flow"]):
                for i, flows in enumerate(node_data["out_flow"]):
                    for flow in flows:
                        #第i个输出对应着节点j的第k个输入
                        j = nodes[flow[0]] #节点j
                        k = flow[1]        #handle k
                        nodes[n].add_outflow(i, j, k)

            print(node_data["options"])
            nodes[n].options = node_data["options"]

        return nodes

imageFlowExecutor = ImageFlowExecutor()